import fs from 'node:fs'
import process from 'node:process'
import path from 'node:path'
import { genStreamingCode as genStreamingCodeFromChatGPT } from 'v1-llm-core'
import { genStreamingCode as genStreamingCodeFromClaude } from 'v1-llm-core/claude'
import type { Observer } from '@trpc/server/observable'
import { observable } from '@trpc/server/observable'
import z from 'zod'
import { v4 as uuidv4 } from 'uuid'
import { type ISearchInput, LLMModel } from '../../../shared/types'
import { publicProcedure } from '../trpc'
import { logger } from '../logger'
import { codeStore } from '../db'

const CODE_BLOCK_SIGN = '```'

function useExtractCodeFromMD() {
  let parseStage: 'init' | 'codeStart' | 'codeBody' | 'complete' = 'init'
  let cache = ''
  // Extract code from streaming markdown
  return (newTrunk: string | null | undefined) => {
    if (!newTrunk)
      return null

    // if receive '```' means the code block start
    if (parseStage === 'init') {
      if (newTrunk.includes(CODE_BLOCK_SIGN)) {
        parseStage = 'codeStart'
        cache += newTrunk.slice(newTrunk.indexOf(CODE_BLOCK_SIGN))
        return null
      }
    }
    else if (parseStage === 'codeStart') {
      cache += newTrunk
      // after received jsx and \n  means the code block body start
      if (cache.includes('```jsx\n')) {
        parseStage = 'codeBody'
        return cache.replace('```jsx\n', '')
      }
      return null
    }
    else if (parseStage === 'codeBody') {
      // receive '```', means the code block end
      if (newTrunk.includes(CODE_BLOCK_SIGN)) {
        parseStage = 'complete'
        return newTrunk.slice(0, newTrunk.indexOf(CODE_BLOCK_SIGN))
      }
      return newTrunk
    }

    return null
  }
}

/**
 * @brief
 * - This trpc router is used to generate code from the question
 * - use observable to stream the result to the client
 * @link
 * https://github.com/trpc/examples-next-prisma-websockets-starter/blob/main/src/server/routers/_app.ts
 */
export const codeGenRouter = publicProcedure
  .input(z.object({
    type: z.string(),
    value: z.string(),
    previousCode: z.string().optional(),
    model: z.string().optional(),
  })) // question
  .subscription(({ input: question }) => {
    return observable<string | null | undefined>((emit) => {
      const id = uuidv4()
      emit.next(`id: ${id}`)
      if (process.env.USE_MOCK)
        mock(emit, id)
      else if (question.model === LLMModel.Claude3)
        genCodeFromClaude(question as ISearchInput, emit, id)
      else
        genCodeFromChatGPT(question as ISearchInput, emit, id)
    })
  })

function genCodeFromChatGPT(question: ISearchInput, emit: Observer<string | null | undefined, unknown>, id: string) {
  (async () => {
    logger.info(`ChatGPT Streaming started for question: ${question}`)
    const extract = useExtractCodeFromMD()
    try {
      const events = await genStreamingCodeFromChatGPT(question)

      let chunks = ''
      let code = ''
      for await (const event of events) {
        // keep streaming until the result is completed
        if (event.choices && event.choices.length > 0) {
          const chunk = event.choices[0].delta?.content
          if (chunk) {
            chunks += chunk
            // send the chunk to the client
            const extractedChunk = extract(chunk)
            if (extractedChunk) {
              emit.next(extractedChunk)
              code += extractedChunk
            }
          }
        }
      }
      logger.info(`Streaming completed for question: ${question}, result: ${chunks}`)
      codeStore.set(id, code)
      emit.complete()
    }
    catch (error) {
      logger.error(`Streaming error for question: ${question}, error: ${error}`)
      emit.error(error)
    }
  })()
}

function genCodeFromClaude(question: ISearchInput, emit: Observer<string | null | undefined, unknown>, id: string) {
  (async () => {
    logger.info(`Claude streaming started for question: ${JSON.stringify(question)}`)
    const extract = useExtractCodeFromMD()
    try {
      let chunks = ''
      let code = ''
      const stream = await genStreamingCodeFromClaude(question)
      stream.on('text', (chunk) => {
        logger.info(chunk)
        if (chunk) {
          chunks += chunk
          // send the chunk to the client
          const extractedChunk = extract(chunk)
          if (extractedChunk) {
            code += extractedChunk
            emit.next(extractedChunk)
          }
        }
      }).on('end', () => {
        logger.info(`Streaming completed for question: ${JSON.stringify(question)}, result: ${chunks}`)
        codeStore.set(id, code)
        emit.complete()
      }).once('error', (error) => {
        logger.error(`Streaming error for question: ${JSON.stringify(question)}, error: ${error}`)
        emit.error(error)
      })
    }
    catch (error) {
      logger.error(`Streaming error for question: ${JSON.stringify(question)}, error: ${error}`)
      emit.error(error)
    }
  })()
}
/**
 * use mock data to streaming
 * @param emit
 */
function mock(emit: Observer<string | null | undefined, unknown>, id: string) {
  logger.info('start mock streaming')
  const extract = useExtractCodeFromMD()

  try {
    const mock = fs.readFileSync(path.join(__dirname, '../../mocks/mock.json')).toString('utf8')
    const mockData = JSON.parse(mock)

    let index = 0
    let code = ''
    function emitWithDelay() {
      if (index < mockData.length) {
        // send the chunk to the client
        const extractedChunk = extract(mockData[index])
        if (extractedChunk) {
          code += extractedChunk
          emit.next(extractedChunk)
        }

        index++
        setTimeout(emitWithDelay, 50) // 每隔 1 秒调用 iterate 函数
      }
      else {
        codeStore.set(id, code)
        emit.complete()
      }
    }

    emitWithDelay()
  }
  catch (e) {
    logger.error(e)
  }
}
